from PIL import Image
import numpy as np
import os
import json

min_thresh = 2 #字符上最少的像素点
min_range = 20 #字符最小的宽度
 
def vertical(img_arr):
    h,w = img_arr.shape
    ver_list = []
    for x in range(w):
        ver_list.append(h - np.count_nonzero(img_arr[:, x]))
    return ver_list

 
def OTSU_enhance(img_gray, th_begin=0, th_end=256, th_step=1):  
    max_g = 0  
    suitable_th = 0  
    for threshold in range(th_begin, th_end, th_step):  
        bin_img = img_gray > threshold  
        bin_img_inv = img_gray <= threshold  
        fore_pix = np.sum(bin_img)  
        back_pix = np.sum(bin_img_inv)  
        if 0 == fore_pix:  
            break  
        if 0 == back_pix:  
            continue  
 
        w0 = float(fore_pix) / img_gray.size  
        u0 = float(np.sum(img_gray * bin_img)) / fore_pix  
        w1 = float(back_pix) / img_gray.size  
        u1 = float(np.sum(img_gray * bin_img_inv)) / back_pix  
        # intra-class variance  
        g = w0 * w1 * (u0 - u1) * (u0 - u1)  
        if g > max_g:  
            max_g = g  
            suitable_th = threshold  
    return suitable_th 
 
 
def simple_cut(vert):
    begin, end = 0,0
    cuts = []
    for i,count in enumerate(vert):
       if count >= min_thresh and begin == 0:
            begin = i
       elif count >= min_thresh and begin != 0:
            continue
       elif count <= min_thresh and begin != 0:
            end = i
            #print (begin, end), count
            if end - begin >= min_range:
                cuts.append((begin, end))
                begin = 0
                end = 0
                continue
       elif count <= min_thresh or begin == 0:
            continue
    return cuts



#根据字数确定刀数
def cut_by_lines(vert,character):
    
    cuts=simple_cut(vert)
    print(cuts)
    if len(cuts)==character:
        return cuts
    else:
        cuts=[]
        p=[]
        v=[]
        for i,count in enumerate(vert):
            if i!=0 and i!=(len(vert)-1):
                if count<=vert[i-1] and count<=vert[i+1]:
                    if count>=10:
                        continue
                    elif count==vert[i-1]:
                        continue
                    else:
                        v.append(count)
                        p.append(i)
        print(p)
        print(v)
        print(vert)
        print(len(vert))

        begin=0
        for x in range(character-1):
            for y in range(len(p)-1):
                if p[y]-begin<10:
                    continue
                else:
                    cuts.append((begin,p[y]))
                    begin=p[y]
        cuts.append((begin,len(vert)-1))
        if len(cuts)==character:
            return cuts
        else:
            cuts=[]
        
            step=len(vert)//character
            begin=0
            for line in range(character-1):
                print(line)
                cuts.append((begin,step*(line+1)))
                begin=step*(line+1)
            cuts.append((begin,len(vert)-1))
            return cuts

save_path = 'test_cut'
 


def cut(pic_path):
    with open("test_num.csv",'r',encoding='utf-8') as f:
        raw_lines=f.readlines()
        f.close()

    result=[]
    for i in raw_lines:
        result.append(i.split('\t'))
    target=os.path.basename(pic_path)
    for line in result:
        if line[0]==target.split('.')[0]:
            character=int(line[1][0])

    src_pic = Image.open(pic_path).convert('L')
    src_arr = np.array(src_pic)
    threshold = OTSU_enhance(src_arr)
    bin_arr = np.where(src_arr < threshold, 0, 255) #先用大津阈值将图片二值化处理
 

    bin_img=Image.fromarray((255 - bin_arr).astype("uint8"))
    width, height = bin_img.size
 
    vert = vertical(bin_arr) #获取到该行的 垂直 方向的投影
    #print(vert)
    #cut = simple_cut(vert) #直接对目标行进行切割
    
    cut = cut_by_lines(vert,character)
    # if len(cut)!=character:
    #     print('exit')
    #     exit(0)
    
 
    for x in range(len(cut)):
        ax = (cut[x][0] - 1, 0, cut[x][1] + 1, height)
        temp = bin_img.crop(ax)
        #temp = image_unit.resize(temp, save_size)
        temp.save('{}/{}_{}.jpg'.format(save_path, os.path.basename(pic_path).split('.')[0],x))




#读取文件目录
def getFileList(dir,Filelist, ext=None):
    newDir = dir
    if os.path.isfile(dir):
        if ext is None:
            Filelist.append(dir)
        else:
            if ext in dir[-3:]:
                Filelist.append(dir)
    
    elif os.path.isdir(dir):
        for s in os.listdir(dir):
            newDir=os.path.join(dir,s)
            getFileList(newDir, Filelist, ext)
 
    return Filelist
 
org_img_folder='./ocr-exercise/test'
 
#检索文件
imglist = getFileList(org_img_folder, [], 'jpg')

 

i=0 
for imgpath in imglist:
    i+=1
    imgname= os.path.splitext(os.path.basename(imgpath))[0]
    
    result=cut(imgpath)
    print('执行到 '+str(i)+' 张图像')